#include "kernel_operator.h"
#include <type_traits>
#include <cmath>

using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
template<typename T> class KernelGather {
public:
    __aicore__ inline KernelGather() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR indices, GM_ADDR y,
        int32_t indices_size, int32_t loopCount, int32_t cycles, int32_t interval,
        uint32_t finalTileNum, uint32_t tileDataNum, uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->batch_cycles = indices_size;
        this->cycles = cycles;
        this->interval = interval;
        this->loopCount = loopCount;

        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;

        x1Gm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X1 *>(x1), 1);
        indicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_INDICES *>(indices), 1);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), 1);
        
        pipe.InitBuffer(inOutQueue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
    }
    __aicore__ inline void Process() {
        int32_t num = 0;
        int32_t increase = this->batch_cycles*this->interval;
        int32_t pian1 = 0;

        {
            for(int64_t i = 0; i < this->loopCount; i++)
            {
                for(int64_t j = 0; j < this->cycles; j++)
                {
                    int64_t indices = indicesGm.GetValue(num);
                    int32_t pian = pian1 + indices * this->interval;

                    int32_t smallloop = this->tileNum;
                    this->processDataNum = this->tileDataNum;
                    for (int32_t z = 0; z < smallloop; z++) {
                        if (z == this->tileNum - 1) {
                        this->processDataNum = this->tailDataNum;
                        }
                        auto bindLocal_in = inOutQueue.AllocTensor<DTYPE_X1>();
                        DataCopy(bindLocal_in, x1Gm[pian + z * this->tileDataNum], this->processDataNum);
                        inOutQueue.EnQue(bindLocal_in);
                        
                        auto bindLocal_out = inOutQueue.DeQue<DTYPE_X1>();
                        DataCopy(yGm[num*this->interval + z * this->tileDataNum], bindLocal_out, this->processDataNum);
                        inOutQueue.FreeTensor(bindLocal_out);
                    }
                    num ++;
                }
                pian1 += increase;
            }
        }
    }
private:
    TPipe pipe;
    // TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> inOutQueue; 
    // TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2, QueueTmp3, QueueTmpBuff;
    TBuf<QuePosition::VECCALC> boox;

    GlobalTensor<DTYPE_INDICES> indicesGm;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_Y> yGm;
    
    int64_t batch_cycles;
    int64_t cycles;
    int64_t interval;
    int64_t loopCount;

    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;

};

template<typename T> class KernelGather_all {
public:
    __aicore__ inline KernelGather_all() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR indices, GM_ADDR y,
        int32_t indices_size, int32_t loopCount, int32_t cycles, int32_t interval) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->batch_cycles = indices_size;
        this->cycles = cycles;
        this->interval = interval;
        this->loopCount = loopCount;

        x1Gm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X1 *>(x1), 1);
        indicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_INDICES *>(indices), 1);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), 1);

    }
    __aicore__ inline void Process() {

        {
            int32_t num = 0;
            int32_t indices_num = 0;
            for(int64_t i = 0; i < this->loopCount; i++)
            {
                int32_t pian1 = i*this->batch_cycles*this->interval;
                for(int64_t j = 0; j < this->cycles; j++)
                {
                    int64_t indices = indicesGm.GetValue(indices_num);
                    indices_num ++;
                    int32_t pian = pian1 + indices * this->interval;
                    for(int64_t z = 0; z < this->interval; z++)
                    {

                        yGm.SetValue(num, x1Gm.GetValue(pian));
                        num ++;
                        pian ++;
                    }
                }
            }
        }
    }
private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> boox;

    GlobalTensor<DTYPE_INDICES> indicesGm;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_Y> yGm;
    
    int64_t batch_cycles;
    int64_t cycles;
    int64_t interval;
    int64_t loopCount;



};
extern "C" __global__ __aicore__ void gather(GM_ADDR x1, GM_ADDR indices, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if (TILING_KEY_IS(1)) {
        KernelGather<DTYPE_X1> op;
        op.Init(x1, indices, y,
                tiling_data.indices_size, tiling_data.loopCount, tiling_data.cycles, tiling_data.interval,
                tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);
        op.Process();
    }
    else if (TILING_KEY_IS(2)) {
        KernelGather_all<DTYPE_X1> op;
        op.Init(x1, indices, y,
                tiling_data.indices_size, tiling_data.loopCount, tiling_data.cycles, tiling_data.interval);
        op.Process();
    }
}